{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp problem_types.contrastive_learning\n",
    "import os\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"-1\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-12 21:42:24.148 | INFO     | m3tl.base_params:register_multiple_problems:526 - Adding new problem weibo_fake_ner, problem type: seq_tag\n",
      "2021-06-12 21:42:24.149 | INFO     | m3tl.base_params:register_multiple_problems:526 - Adding new problem weibo_fake_multi_cls, problem type: multi_cls\n",
      "2021-06-12 21:42:24.149 | INFO     | m3tl.base_params:register_multiple_problems:526 - Adding new problem weibo_fake_cls, problem type: cls\n",
      "2021-06-12 21:42:24.150 | INFO     | m3tl.base_params:register_multiple_problems:526 - Adding new problem weibo_masklm, problem type: masklm\n",
      "2021-06-12 21:42:24.150 | INFO     | m3tl.base_params:register_multiple_problems:526 - Adding new problem weibo_fake_regression, problem type: regression\n",
      "2021-06-12 21:42:24.151 | INFO     | m3tl.base_params:register_multiple_problems:526 - Adding new problem weibo_fake_vector_fit, problem type: vector_fit\n",
      "2021-06-12 21:42:24.151 | INFO     | m3tl.base_params:register_multiple_problems:526 - Adding new problem weibo_premask_mlm, problem type: premask_mlm\n",
      "2021-06-12 21:42:24.152 | INFO     | m3tl.base_params:register_multiple_problems:526 - Adding new problem fake_contrastive_learning, problem type: contrastive_learning\n",
      "2021-06-12 21:42:24.152 | WARNING  | m3tl.base_params:assign_problem:620 - base_dir and dir_name arguments will be deprecated in the future. Please use model_dir instead.\n",
      "2021-06-12 21:42:24.153 | WARNING  | m3tl.base_params:prepare_dir:364 - bert_config not exists. will load model from huggingface checkpoint.\n",
      "2021-06-12 21:42:30.466 | WARNING  | m3tl.read_write_tfrecord:chain_processed_data:248 - Chaining problems with & may consume a lot of memory if data is not pyspark RDD.\n",
      "2021-06-12 21:42:30.481 | DEBUG    | m3tl.read_write_tfrecord:_write_fn:134 - Writing /tmp/tmpyjjr_dm7/weibo_fake_cls_weibo_fake_ner_weibo_fake_regression_weibo_fake_vector_fit/train_00000.tfrecord\n",
      "2021-06-12 21:42:30.547 | WARNING  | m3tl.read_write_tfrecord:chain_processed_data:248 - Chaining problems with & may consume a lot of memory if data is not pyspark RDD.\n",
      "2021-06-12 21:42:30.561 | DEBUG    | m3tl.read_write_tfrecord:_write_fn:134 - Writing /tmp/tmpyjjr_dm7/weibo_fake_cls_weibo_fake_ner_weibo_fake_regression_weibo_fake_vector_fit/eval_00000.tfrecord\n",
      "2021-06-12 21:42:30.588 | DEBUG    | m3tl.read_write_tfrecord:_write_fn:134 - Writing /tmp/tmpyjjr_dm7/weibo_fake_multi_cls/train_00000.tfrecord\n",
      "2021-06-12 21:42:30.612 | DEBUG    | m3tl.read_write_tfrecord:_write_fn:134 - Writing /tmp/tmpyjjr_dm7/weibo_fake_multi_cls/eval_00000.tfrecord\n",
      "2021-06-12 21:42:30.688 | DEBUG    | m3tl.read_write_tfrecord:_write_fn:134 - Writing /tmp/tmpyjjr_dm7/weibo_masklm/train_00000.tfrecord\n",
      "2021-06-12 21:42:30.736 | DEBUG    | m3tl.read_write_tfrecord:_write_fn:134 - Writing /tmp/tmpyjjr_dm7/weibo_masklm/eval_00000.tfrecord\n",
      "2021-06-12 21:42:30.801 | DEBUG    | m3tl.read_write_tfrecord:_write_fn:134 - Writing /tmp/tmpyjjr_dm7/weibo_premask_mlm/train_00000.tfrecord\n",
      "2021-06-12 21:42:30.864 | DEBUG    | m3tl.read_write_tfrecord:_write_fn:134 - Writing /tmp/tmpyjjr_dm7/weibo_premask_mlm/eval_00000.tfrecord\n",
      "2021-06-12 21:42:30.883 | DEBUG    | m3tl.read_write_tfrecord:_write_fn:134 - Writing /tmp/tmpyjjr_dm7/fake_contrastive_learning/train_00000.tfrecord\n",
      "2021-06-12 21:42:30.900 | DEBUG    | m3tl.read_write_tfrecord:_write_fn:134 - Writing /tmp/tmpyjjr_dm7/fake_contrastive_learning/eval_00000.tfrecord\n",
      "2021-06-12 21:42:31.923 | INFO     | m3tl.input_fn:train_eval_input_fn:56 - sampling weights: \n",
      "2021-06-12 21:42:31.924 | INFO     | m3tl.input_fn:train_eval_input_fn:57 - {\n",
      "    \"weibo_fake_cls_weibo_fake_ner_weibo_fake_regression_weibo_fake_vector_fit\": 0.2127659574468085,\n",
      "    \"weibo_fake_multi_cls\": 0.2127659574468085,\n",
      "    \"weibo_masklm\": 0.14893617021276595,\n",
      "    \"weibo_premask_mlm\": 0.2127659574468085,\n",
      "    \"fake_contrastive_learning\": 0.2127659574468085\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "# test setup\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from m3tl.input_fn import train_eval_input_fn\n",
    "from m3tl.test_base import TestBase, test_top_layer\n",
    "\n",
    "test_base = TestBase()\n",
    "params = test_base.params\n",
    "\n",
    "hidden_dim = params.bert_config.hidden_size\n",
    "\n",
    "train_dataset = train_eval_input_fn(params=params)\n",
    "one_batch = next(train_dataset.as_numpy_iterator())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(32, 1, 10)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "one_batch['array_input_ids'].shape\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Contrastive Learning(contrastive_learning)\n",
    "\n",
    "Contrastive learning is usually used along side in-batch data augmentation. To implement data augmentation strategy, one should implement their own embedding layer. See `embedding_layer` sub-module for more details.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports and utils\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "from typing import List\n",
    "\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from loguru import logger\n",
    "from m3tl.base_params import BaseParams\n",
    "from m3tl.problem_types.utils import (empty_tensor_handling_loss,\n",
    "                                      nan_loss_handling, pad_to_shape)\n",
    "from m3tl.special_tokens import PREDICT\n",
    "from m3tl.utils import (LabelEncoder, get_label_encoder_save_path, get_phase,\n",
    "                        need_make_label_encoder)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Top Layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SimSCE\n",
    "# export\n",
    "class SimCSE(tf.keras.Model):\n",
    "    def __init__(self, params: BaseParams, problem_name: str) -> None:\n",
    "        super(SimCSE, self).__init__(name='simcse')\n",
    "        self.params = params\n",
    "        self.problem_name = problem_name\n",
    "        self.dropout = tf.keras.layers.Dropout(self.params.dropout)\n",
    "        self.pooler = self.params.get('simcse_pooler', 'pooled')\n",
    "        self.metric_fn = tf.keras.metrics.CategoricalAccuracy(name='{}_acc'.format(problem_name))\n",
    "        availabel_pooler = ['pooled', 'mean_pool']\n",
    "        assert self.pooler in availabel_pooler, \\\n",
    "            'available params.simcse_pooler: {}, got: {}'.format(\n",
    "                availabel_pooler, self.pooler)\n",
    "        if self.params.embedding_layer['name'] != 'duplicate_data_augmentation_embedding':\n",
    "            raise ValueError(\n",
    "                'SimCSE requires duplicate_data_augmentation_embedding. Fix it with `params.assign_embedding_layer(\\'duplicate_data_augmentation_embedding\\')`')\n",
    "\n",
    "    def call(self, inputs):\n",
    "\n",
    "        features, hidden_features = inputs\n",
    "        phase = get_phase()\n",
    "\n",
    "        if phase != PREDICT:\n",
    "            # created pool embedding\n",
    "            if self.pooler == 'pooled':\n",
    "                all_pooled_embedding = hidden_features['pooled']\n",
    "            else:\n",
    "                all_pooled_embedding = tf.reduce_mean(\n",
    "                    hidden_features['seq'], axis=1)\n",
    "\n",
    "            # shape (batch_size, hidden_dim)\n",
    "            pooled_rep1_embedding, pooled_rep2_embedding = tf.split(\n",
    "                all_pooled_embedding, 2)\n",
    "\n",
    "            # calculate similarity\n",
    "            pooled_rep1_embedding = tf.math.l2_normalize(\n",
    "                pooled_rep1_embedding, axis=1)\n",
    "            pooled_rep2_embedding = tf.math.l2_normalize(\n",
    "                pooled_rep2_embedding, axis=1)\n",
    "            # shape (batch_size, batch_size)\n",
    "            similarity = tf.matmul(pooled_rep1_embedding,\n",
    "                                   pooled_rep2_embedding, transpose_b=True)\n",
    "            labels = tf.eye(tf.shape(similarity)[0])\n",
    "\n",
    "            # shape (batch_size*batch_size)\n",
    "            similarity = tf.reshape(similarity, shape=(-1, 1))\n",
    "            labels = tf.reshape(labels, shape=(-1, 1))\n",
    "\n",
    "            # make compatible with binary crossentropy\n",
    "            similarity = tf.concat([1-similarity, similarity], axis=1)\n",
    "            labels = tf.concat([1-labels, labels], axis=1)\n",
    "            loss = tf.keras.losses.binary_crossentropy(labels, similarity)\n",
    "            loss = tf.reduce_mean(loss)\n",
    "            self.add_loss(loss)\n",
    "            acc = self.metric_fn(labels, similarity)\n",
    "            self.add_metric(acc)\n",
    "        return inputs[1]['pooled']\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def get_contrastive_learning_model(params: BaseParams, problem_name: str, model_name: str) -> tf.keras.Model:\n",
    "    if model_name == 'simcse':\n",
    "        return SimCSE(params=params, problem_name=problem_name)\n",
    "\n",
    "    logger.warning(\n",
    "        '{} not match any contrastive learning model, using SimCSE'.format(model_name))\n",
    "    return SimCSE(params=params, problem_name=problem_name)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "\n",
    "class ContrastiveLearning(tf.keras.Model):\n",
    "    def __init__(self, params: BaseParams, problem_name: str) -> None:\n",
    "        super(ContrastiveLearning, self).__init__(name=problem_name)\n",
    "        self.params = params\n",
    "        self.problem_name = problem_name\n",
    "        self.contrastive_learning_model_name = self.params.contrastive_learning_model_name\n",
    "        self.contrastive_learning_model = get_contrastive_learning_model(\n",
    "            params=self.params, problem_name=problem_name, model_name=self.contrastive_learning_model_name)\n",
    "\n",
    "    def call(self, inputs):\n",
    "        return self.contrastive_learning_model(inputs)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-12 21:42:33.553 | DEBUG    | m3tl.test_base:test_top_layer:247 - Testing ContrastiveLearning\n",
      "2021-06-12 21:42:33.569 | WARNING  | __main__:get_contrastive_learning_model:7 - None not match any contrastive learning model, using SimCSE\n",
      "2021-06-12 21:42:33.585 | DEBUG    | m3tl.test_base:test_top_layer:253 - testing batch size 0\n",
      "2021-06-12 21:42:33.593 | DEBUG    | m3tl.test_base:test_top_layer:253 - testing batch size 2\n"
     ]
    }
   ],
   "source": [
    "test_top_layer(ContrastiveLearning, problem='fake_contrastive_learning',\n",
    "               params=params, sample_features=one_batch, hidden_dim=hidden_dim, test_batch_size_list=[0,2])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get or make label encoder function\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def contrastive_learning_get_or_make_label_encoder_fn(params: BaseParams, problem: str, mode: str, label_list: List[str], *args, **kwargs) -> LabelEncoder:\n",
    "\n",
    "    le_path = get_label_encoder_save_path(params=params, problem=problem)\n",
    "    label_encoder = LabelEncoder()\n",
    "    if need_make_label_encoder(mode=mode, le_path=le_path, overwrite=kwargs['overwrite']):\n",
    "        # fit and save label encoder\n",
    "        label_encoder.fit(label_list)\n",
    "        label_encoder.dump(le_path)\n",
    "        params.set_problem_info(\n",
    "            problem=problem, info_name='num_classes', info=len(label_encoder.encode_dict))\n",
    "    else:\n",
    "        label_encoder.load(le_path)\n",
    "\n",
    "    return label_encoder\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Label handing function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def contrastive_learning_label_handling_fn(target: str, label_encoder=None, tokenizer=None, decoding_length=None, *args, **kwargs) -> dict:\n",
    "\n",
    "    label_id = label_encoder.transform([target]).tolist()[0]\n",
    "    label_id = np.int32(label_id)\n",
    "    return label_id, None\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.10 64-bit ('base': conda)",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
